"""
外汇数据采集器
"""
import logging
import pandas as pd
import datetime
import akshare as ak
from sqlalchemy.exc import IntegrityError
from ..models.base import get_db
from ..models.financial import ForexData, ForexDailyData

logger = logging.getLogger(__name__)

class ForexCollector:
    """外汇数据采集器"""
    
    def __init__(self):
        self.db = next(get_db())
        self.batch_size = 100
        
    def collect_forex_info(self, currency_type="外汇"):
        """
        采集外汇基本信息
        
        Args:
            currency_type: 货币类型(外汇/加密货币)
            
        Returns:
            (新增数量, 更新数量)
        """
        try:
            # 调用AKShare接口获取外汇信息
            df = ak.forex_info(currency_type=currency_type)
            
            if df.empty:
                logger.info(f"没有可用的{currency_type}信息")
                return 0, 0
                
            # 数据清洗
            df = self._clean_forex_info(df)
            
            # 保存到数据库
            new_count, update_count = self._save_forex_info(df)
            
            logger.info(f"成功保存{currency_type}信息: 新增{new_count}条, 更新{update_count}条")
            return new_count, update_count
            
        except Exception as e:
            logger.error(f"采集{currency_type}信息失败: {str(e)}")
            raise
            
    def collect_forex_daily(self, symbol, start_date=None, end_date=None, task_id=None):
        """
        采集外汇日线数据
        
        Args:
            symbol: 外汇代码(如USD/CNY)
            start_date: 开始日期(YYYYMMDD)
            end_date: 结束日期(YYYYMMDD)
            task_id: 任务ID
            
        Returns:
            成功保存的数据条数
        """
        # 更新任务状态
        if task_id:
            logger.info(f"开始执行任务 {task_id}: 采集外汇{symbol}日线数据")
            
        try:
            # 调用AKShare接口获取外汇日线数据
            df = ak.forex_daily(symbol=symbol, start_date=start_date, end_date=end_date)
            
            if df.empty:
                logger.info(f"外汇{symbol}没有可用的日线数据")
                if task_id:
                    logger.info(f"任务 {task_id} 完成: 没有可用数据")
                return 0
                
            # 数据清洗
            df = self._clean_forex_daily(df)
            
            # 保存到数据库
            inserted_count = self._save_forex_daily(symbol, df)
            
            # 更新任务状态
            if task_id:
                logger.info(f"任务 {task_id} 完成: 成功保存 {inserted_count} 条数据")
                
            return inserted_count
            
        except Exception as e:
            error_msg = f"采集外汇{symbol}日线数据失败: {str(e)}"
            logger.error(error_msg)
            if task_id:
                logger.error(f"任务 {task_id} 失败: {error_msg}")
            raise
            
    def _clean_forex_info(self, df):
        """清洗外汇基本信息"""
        # 标准化列名
        df.columns = [col.strip() for col in df.columns]
        
        # 处理日期格式
        if '上市日期' in df.columns:
            df['上市日期'] = pd.to_datetime(df['上市日期'])
            
        return df
        
    def _save_forex_info(self, df):
        """保存外汇信息到数据库"""
        new_count = 0
        update_count = 0
        
        try:
            for _, row in df.iterrows():
                # 检查是否已存在
                existing = self.db.query(ForexData).filter(
                    ForexData.symbol == row['货币代码']
                ).first()
                
                if existing:
                    # 更新现有记录
                    existing.name = row['货币名称']
                    existing.currency_type = row['货币类型']
                    existing.update_time = datetime.datetime.now()
                    update_count += 1
                else:
                    # 新增记录
                    forex = ForexData(
                        symbol=row['货币代码'],
                        name=row['货币名称'],
                        currency_type=row['货币类型'],
                        update_time=datetime.datetime.now()
                    )
                    self.db.add(forex)
                    new_count += 1
                    
                # 分批提交
                if (new_count + update_count) % self.batch_size == 0:
                    self.db.commit()
                    
            self.db.commit()
            return new_count, update_count
            
        except IntegrityError as e:
            self.db.rollback()
            logger.error(f"保存外汇信息失败(完整性错误): {str(e)}")
            raise
        except Exception as e:
            self.db.rollback()
            logger.error(f"保存外汇信息失败: {str(e)}")
            raise
            
    def _clean_forex_daily(self, df):
        """清洗外汇日线数据"""
        # 标准化列名
        df.columns = [col.strip() for col in df.columns]
        
        # 处理日期格式
        if '日期' in df.columns:
            df['日期'] = pd.to_datetime(df['日期'])
            
        return df
        
    def _save_forex_daily(self, symbol, df):
        """保存外汇日线数据到数据库"""
        inserted_count = 0
        
        try:
            # 获取外汇ID
            forex = self.db.query(ForexData).filter(
                ForexData.symbol == symbol
            ).first()
            
            if not forex:
                raise ValueError(f"外汇{symbol}不存在")
                
            # 准备批量插入
            batch_records = []
            
            for _, row in df.iterrows():
                forex_daily = ForexDailyData(
                    forex_id=forex.id,
                    date=row['日期'],
                    open=row['开盘价'],
                    high=row['最高价'],
                    low=row['最低价'],
                    close=row['收盘价'],
                    update_time=datetime.datetime.now()
                )
                batch_records.append(forex_daily)
                inserted_count += 1
                
                # 分批提交
                if len(batch_records) >= self.batch_size:
                    self.db.add_all(batch_records)
                    self.db.commit()
                    batch_records = []
                    
            # 提交剩余记录
            if batch_records:
                self.db.add_all(batch_records)
                self.db.commit()
                
            return inserted_count
            
        except IntegrityError as e:
            self.db.rollback()
            logger.error(f"保存外汇日线数据失败(完整性错误): {str(e)}")
            raise
        except Exception as e:
            self.db.rollback()
            logger.error(f"保存外汇日线数据失败: {str(e)}")
            raise